package main

import (
	"fmt"
	"math"
	"sync"
	"time"
)

// 生成器函数：生成整数并发送到返回的channel上
func generator(nums ...int) <-chan int {
	out := make(chan int)

	go func() {
		defer close(out) // 确保通道最终会被关闭

		for _, n := range nums {
			out <- n // 将数据发送到通道
		}
	}()

	return out
}

// 平方函数：接收一个整数通道，返回每个整数平方后的通道
func square(in <-chan int) <-chan int {
	out := make(chan int)

	go func() {
		defer close(out)

		for n := range in {
			out <- n * n                       // 对输入的每个数字求平方并输出
			time.Sleep(100 * time.Millisecond) // 模拟一些处理时间
		}
	}()

	return out
}

// 筛选函数：保留能被3整除的数字
func filter(in <-chan int) <-chan int {
	out := make(chan int)

	go func() {
		defer close(out)

		for n := range in {
			if n%3 == 0 { // 只保留能被3整除的数字
				out <- n
			}
			time.Sleep(50 * time.Millisecond) // 模拟一些处理时间
		}
	}()

	return out
}

// 求和函数：计算通道中所有整数的和
func sum(in <-chan int) <-chan int {
	out := make(chan int)

	go func() {
		defer close(out)

		total := 0
		for n := range in {
			total += n
		}

		out <- total
	}()

	return out
}

// 计算平方根函数：接收一个整数通道，返回每个整数平方根后的通道
func sqrt(in <-chan int) <-chan float64 {
	out := make(chan float64)

	go func() {
		defer close(out)

		for n := range in {
			out <- math.Sqrt(float64(n))
			time.Sleep(80 * time.Millisecond) // 模拟一些处理时间
		}
	}()

	return out
}

// 合并多个通道的输出到一个通道
func merge(cs ...<-chan int) <-chan int {
	var wg sync.WaitGroup
	out := make(chan int)

	// 为每个输入通道启动一个goroutine
	output := func(c <-chan int) {
		defer wg.Done()
		for n := range c {
			out <- n
		}
	}

	wg.Add(len(cs))
	for _, c := range cs {
		go output(c)
	}

	// 启动一个goroutine，在所有输入关闭后关闭输出通道
	go func() {
		wg.Wait()
		close(out)
	}()

	return out
}

// 扇出：同一个数据源被多个接收者并发处理
// 扇入：多个结果被合并到一个通道
func fanOutFanIn(in <-chan int) <-chan int {
	// 扇出到3个square处理器
	c1 := square(in)
	c2 := square(in)
	c3 := square(in)

	// 扇入：合并3个处理器的结果
	return merge(c1, c2, c3)
}

func main() {
	// 创建管道：生成器 -> 平方 -> 求和
	fmt.Println("简单的pipeline示例：生成器 -> 平方 -> 求和")
	nums := []int{2, 3, 4, 5, 6, 7, 8, 9, 10}

	// 生成一个包含1到10的整数通道
	fmt.Println("输入数据:", nums)
	numChan := generator(nums...)

	// 计算每个数的平方
	squareChan := square(numChan)

	// 求平方和
	sumChan := sum(squareChan)

	// 输出结果
	fmt.Println("平方和:", <-sumChan)

	// 更复杂的管道示例
	fmt.Println("复杂的pipeline示例：生成器 -> 平方 -> 筛选 -> 求和")
	nums2 := []int{1, 2, 3, 4, 5, 6, 7, 8, 9}
	fmt.Println("输入数据:", nums2)

	gen := generator(nums2...)
	sq := square(gen)
	filtered := filter(sq)
	sumResult := sum(filtered)

	fmt.Println("被3整除的平方和:", <-sumResult)

	// 扇出-扇入模式示例
	fmt.Println("扇出-扇入模式示例：")
	nums3 := []int{1, 2, 3, 4, 5}
	fmt.Println("输入数据:", nums3)

	start := time.Now()

	gen3 := generator(nums3...)

	// 单处理器管道
	singleStart := time.Now()
	singleResult := sum(square(gen3))
	fmt.Printf("单处理器结果: %d, 耗时: %v\n", <-singleResult, time.Since(singleStart))

	// 多处理器并行管道（扇出-扇入）
	gen4 := generator(nums3...)
	fanStart := time.Now()
	fanResult := sum(fanOutFanIn(gen4))
	fmt.Printf("扇出-扇入结果: %d, 耗时: %v\n", <-fanResult, time.Since(fanStart))

	fmt.Printf("总耗时: %v\n", time.Since(start))

	// 转换管道示例：整数 -> 平方 -> 平方根
	fmt.Println("转换管道示例：整数 -> 平方 -> 平方根")
	nums4 := []int{1, 4, 9, 16, 25}
	fmt.Println("输入数据:", nums4)

	gen5 := generator(nums4...)
	squares := square(gen5)
	roots := sqrt(squares)

	// 收集并打印结果
	var results []float64
	for r := range roots {
		results = append(results, r)
	}
	fmt.Println("平方根结果:", results)
}
